#include <iostream>
#include <fstream>
#include <vector>
#include <cstdint>
#include "acl/acl.h"
#include "acl/ops/acl_cblas.h"

#define CHECK_ACL(expr) \
    do { \
        aclError ret = (expr); \
        if (ret != ACL_SUCCESS) { \
            std::cerr << "ACL Error [" << __FILE__ << ":" << __LINE__ << "] Code=" << ret << std::endl; \
            return EXIT_FAILURE; \
        } \
    } while (0)

// 加载二进制文件数据
std::vector<char> LoadBinaryFile(const std::string& file_path) {
    std::ifstream file(file_path, std::ios::binary | std::ios::ate);
    if (!file.is_open()) {
        throw std::runtime_error("Failed to open file: " + file_path);
    }
    size_t size = file.tellg();
    file.seekg(0, std::ios::beg);
    std::vector<char> buffer(size);
    file.read(buffer.data(), size);
    return buffer;
}

// 保存二进制文件数据
void SaveBinaryFile(const std::string& file_path, const void* data, size_t size) {
    std::ofstream file(file_path, std::ios::binary);
    if (!file.is_open()) {
        throw std::runtime_error("Failed to create file: " + file_path);
    }
    file.write(reinterpret_cast<const char*>(data), size);
}

int main(int argc, char** argv) {
    // 参数解析
    if (argc != 5) {
        std::cerr << "Usage: " << argv[0] << " --input <input.bin> --output <output.bin> --shape <N> --dtype <float32|float16|double>" << std::endl;
        return EXIT_FAILURE;
    }

    std::string input_path = argv[1];
    std::string output_path = argv[2];
    int64_t shape = std::stoll(argv[3]);
    aclDataType dtype = (strcmp(argv[4], "float16") == 0) ? ACL_FLOAT16 :
                       (strcmp(argv[4], "double") == 0) ? ACL_DOUBLE : ACL_FLOAT;

    // 初始化ACL
    CHECK_ACL(aclInit(nullptr));
    CHECK_ACL(aclrtSetDevice(0));

    // 加载输入数据
    auto input_data = LoadBinaryFile(input_path);
    void* dev_input = nullptr;
    CHECK_ACL(aclrtMalloc(&dev_input, input_data.size(), ACL_MEM_MALLOC_NORMAL_ONLY));
    CHECK_ACL(aclrtMemcpy(dev_input, input_data.size(), input_data.data(), input_data.size(), ACL_MEMCPY_HOST_TO_DEVICE));

    // 准备输出内存
    int32_t host_output = 0;
    void* dev_output = nullptr;
    CHECK_ACL(aclrtMalloc(&dev_output, sizeof(int32_t), ACL_MEM_MALLOC_NORMAL_ONLY));

    // 创建算子描述
    aclTensorDesc* input_desc = aclCreateTensorDesc(ACL_FLOAT, 1, &shape, ACL_FORMAT_ND);
    aclTensorDesc* output_desc = aclCreateTensorDesc(ACL_INT32, 1, &shape, ACL_FORMAT_ND);

    // 调用isamin算子
    CHECK_ACL(aclopExecute("Isamin", 1, &input_desc, &dev_input, 
                           1, &output_desc, &dev_output, 
                           nullptr, nullptr, nullptr));

    // 拷贝结果回主机
    CHECK_ACL(aclrtMemcpy(&host_output, sizeof(int32_t), dev_output, sizeof(int32_t), ACL_MEMCPY_DEVICE_TO_HOST));

    // 保存结果
    SaveBinaryFile(output_path, &host_output, sizeof(int32_t));

    // 打印结果 (1-based索引)
    std::cout << "isamin result (1-based index): " << host_output << std::endl;

    // 释放资源
    aclDestroyTensorDesc(input_desc);
    aclDestroyTensorDesc(output_desc);
    aclrtFree(dev_input);
    aclrtFree(dev_output);
    aclrtResetDevice(0);
    aclFinalize();

    return EXIT_SUCCESS;
}